#!/usr/bin/env python3
"""
修正版本的 MCP 客户端
遵循正确的 MCP 协议格式
"""

import asyncio
import aiohttp
import json
import logging
from typing import Dict, Any, Optional, List
from datetime import datetime
import openai
import os

class CorrectedMCPClient:
    """修正版本的 MCP 客户端"""
    
    def __init__(self, mcp_server_url: str = "http://localhost:3000"):
        self.mcp_server_url = mcp_server_url
        self.session = None
        self.logger = logging.getLogger(__name__)
        self.request_id = 1
        
    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()
    
    def get_next_id(self) -> int:
        """获取下一个请求 ID"""
        current_id = self.request_id
        self.request_id += 1
        return current_id
    
    async def list_tools(self) -> Dict[str, Any]:
        """列出可用工具（标准 MCP 消息）"""
        message = {
            "jsonrpc": "2.0",
            "id": self.get_next_id(),
            "method": "tools/list",
            "params": {
                "cursor": None
            }
        }
        
        return await self._send_mcp_message(message)
    
    async def call_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """调用工具（标准 MCP 消息）"""
        message = {
            "jsonrpc": "2.0", 
            "id": self.get_next_id(),
            "method": "tools/call",
            "params": {
                "name": tool_name,
                "arguments": arguments
            }
        }
        
        self.logger.info(f"🔧 调用 MCP 工具: {tool_name}")
        self.logger.debug(f"参数: {arguments}")
        
        return await self._send_mcp_message(message)
    
    async def _send_mcp_message(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """发送 MCP 消息"""
        try:
            # 对于 Playwright MCP，我们需要使用正确的端点
            # 实际上，Playwright MCP 可能不是通过 HTTP POST 进行通信
            # 而是通过 SSE (Server-Sent Events) 连接
            
            self.logger.debug(f"发送 MCP 消息: {message}")
            
            async with self.session.post(
                f"{self.mcp_server_url}/messages",
                json=message,
                headers={
                    "Content-Type": "application/json",
                    "Accept": "application/json"  # 修正：不一定是 SSE
                }
            ) as resp:
                
                if resp.status == 200:
                    result = await resp.json()
                    self.logger.debug(f"收到 MCP 响应: {result}")
                    return result
                else:
                    error_text = await resp.text()
                    self.logger.error(f"❌ MCP 消息发送失败: {resp.status} - {error_text}")
                    return {
                        "jsonrpc": "2.0",
                        "id": message.get("id"),
                        "error": {
                            "code": resp.status,
                            "message": error_text
                        }
                    }
                    
        except Exception as e:
            self.logger.error(f"❌ MCP 消息发送异常: {e}")
            return {
                "jsonrpc": "2.0", 
                "id": message.get("id"),
                "error": {
                    "code": -32603,
                    "message": f"Internal error: {str(e)}"
                }
            }

class PlaywrightMCPDemo:
    """Playwright MCP 演示"""
    
    def __init__(self, mcp_client: CorrectedMCPClient):
        self.client = mcp_client
        self.logger = logging.getLogger(__name__)
    
    async def demo_workflow(self):
        """演示工作流"""
        try:
            # 1. 首先列出可用工具
            self.logger.info("📋 获取可用工具列表...")
            tools_response = await self.client.list_tools()
            
            if "error" in tools_response:
                self.logger.error(f"获取工具失败: {tools_response['error']}")
                return False
            
            available_tools = tools_response.get("result", {}).get("tools", [])
            self.logger.info(f"✅ 可用工具数量: {len(available_tools)}")
            
            for tool in available_tools:
                self.logger.info(f"  - {tool.get('name')}: {tool.get('description')}")
            
            # 2. 调用浏览器导航工具
            nav_result = await self.client.call_tool(
                "mcp_Playwright_browser_navigate",
                {"url": "https://www.xiaohongshu.com"}
            )
            
            if "error" in nav_result:
                self.logger.error(f"导航失败: {nav_result['error']}")
                return False
            
            self.logger.info("✅ 导航成功")
            
            # 3. 获取页面快照
            snapshot_result = await self.client.call_tool(
                "mcp_Playwright_browser_snapshot", 
                {"random_string": "demo"}
            )
            
            if "error" in snapshot_result:
                self.logger.error(f"获取快照失败: {snapshot_result['error']}")
                return False
            
            self.logger.info("✅ 获取页面快照成功")
            
            return True
            
        except Exception as e:
            self.logger.error(f"❌ 演示工作流失败: {e}")
            return False

async def test_real_mcp_connection():
    """测试真实的 MCP 连接"""
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
    logger = logging.getLogger(__name__)
    
    # 首先检查 MCP 服务器是否运行
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get("http://localhost:3000/health") as resp:
                if resp.status != 200:
                    logger.error("❌ MCP 服务器未运行，请先启动服务器")
                    logger.info("💡 启动命令: node playwright-mcp-server.js")
                    return
                else:
                    health_data = await resp.json()
                    logger.info(f"✅ MCP 服务器健康: {health_data}")
    except Exception as e:
        logger.error(f"❌ 无法连接到 MCP 服务器: {e}")
        logger.info("💡 请确保服务器正在运行: node playwright-mcp-server.js")
        return
    
    # 测试 MCP 通信
    async with CorrectedMCPClient() as client:
        demo = PlaywrightMCPDemo(client)
        success = await demo.demo_workflow()
        
        if success:
            logger.info("🎉 MCP 演示成功完成！")
        else:
            logger.error("❌ MCP 演示失败")

if __name__ == "__main__":
    asyncio.run(test_real_mcp_connection())
